import time
from contextlib import contextmanager
from psycopg2.extras import execute_batch
import psycopg2
from psycopg2.pool import ThreadedConnectionPool
from logging_config import logger
from config import read_config

jsonData = read_config()
# 获取数据库配置
database_config = jsonData['db_config']
dbname = database_config['dbname']
user = database_config['user']
password = database_config['password']
host = database_config['host']
port = database_config['port']
min_conn = database_config['min_conn']
max_conn = database_config['max_conn']

database_pool = ThreadedConnectionPool(
    minconn=min_conn, maxconn=max_conn, dbname=dbname, user=user,
    password=password, host=host, port=port
)
logger.info(f'已成功创建PostgreSQL连接池 {database_pool}')


@contextmanager
def get_cursor():
    conn = database_pool.getconn()
    conn.autocommit = True
    try:
        yield conn.cursor()
    finally:
        database_pool.putconn(conn)


def batch_insert(results, table_name, columns):
    with get_cursor() as cursor:
        placeholders = ', '.join(['%s'] * len(columns))
        insert_sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({placeholders})"
        values_list = [list(item[col] for col in columns) for item in results]
        try:
            start = time.time()
            # 执行批量插入
            execute_batch(cursor, insert_sql, values_list)
        except psycopg2.Error as e:
            logger.error(f'插入数据时出错: {e}')
            cursor.rollback()
        finally:
            end = time.time()
            logger.info(f'{table_name} {len(results)}条, 耗时: 总耗时: {end - start :.6f}秒 插入完成')


def insert(result, table_name, columns):
    with get_cursor() as cursor:
        values = [result[col] for col in columns]
        placeholders = ', '.join(['%s'] * len(columns))
        insert_sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({placeholders})"
        try:
            start = time.time()
            cursor.execute(insert_sql, values)
        except psycopg2.Error as e:
            logger.error(f'插入数据时出错: {e}')
            cursor.rollback()
        finally:
            end = time.time()
            logger.info(f'{table_name} {cursor.rowcount}条, 耗时: 总耗时: {end - start :.6f}秒 插入完成')
